# -*- coding: utf-8 -*-

import os
import pandas as pd


def load_data(data_file:str, columns=None, **kwargs):
    """使用pandas库加载指定数据集文件
    
    Arguments:
        data_file {str} -- 文件的路径
    
    Keyword Arguments:
        columns {list} -- 要使用的列 (default: {[]})
    
    Returns:
        pandas.DataFrame
    """

    assert os.path.exists(data_file), 'File %s not found' % data_file
    assert os.path.isfile(data_file), '%s is not a valid file' % data_file
    if columns:
        data = pd.read_csv(data_file, usecols=columns, **kwargs)
    else:
        data = pd.read_csv(data_file, **kwargs)
    return data


def split_dataset(datafilepath:str, p=0.2):
    """分割训练集和测试集
    
    Arguments:
        datafilepath {str} -- 原始数据集路径
    
    Keyword Arguments:
        p {float} -- 分割比例 (default: {0.2})
    """
    dirname = os.path.dirname(datafilepath)
    filename = os.path.splitext(os.path.basename(datafilepath))[0]
    trainfile = os.path.join(dirname, filename+'-train.csv')
    testfile = os.path.join(dirname, filename+'-test.csv')
    datafile:pd.DataFrame = pd.read_csv(datafilepath)
    rows = datafile.shape[0]
    mid = int(rows * p)
    traindata: pd.DataFrame = datafile[0: rows-mid+1]
    testdata: pd.DataFrame = datafile[-mid:]
    traindata.to_csv(trainfile, index=None)
    testdata.to_csv(testfile, index=None)
    print('分割训练集保存在: %s' % trainfile)
    print('分割测试集保存在: %s' % testfile)

